import  torch
from torch import nn, optim
from torch.nn import functional as F
import torchvision
from    torch.utils.data import DataLoader
from    torchvision import transforms


class MyRNN(nn.Module):
    def __init__(self, word_dim, h_dim):
        super(MyRNN, self).__init__()
        self.h_dim = h_dim
        self.cell0 = nn.RNNCell(word_dim, h_dim)
        self.cell1 = nn.RNNCell(h_dim, h_dim)

        self.outlayer = nn.Linear(h_dim, 1)

    def forward(self, inputs):

        for xt in inputs.split(1, dim=0):
            xt = xt.squeeze(0)
            h0 = self.cell0(xt, torch.zeros(inputs.shape[1], self.h_dim))
            h1 = self.cell1(h0, torch.zeros(inputs.shape[1], self.h_dim))
        out = self.outlayer(h1)

        prob = F.sigmoid(out)
        return prob

data = torch.randn(20, 128, 100)
model = MyRNN(100, 20)

out = model(data)
print(out.shape)
